package discovery

import (
	"context"
	"dealcenter/utils/logx"
	"encoding/json"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
	"go.etcd.io/etcd/client/v3"
	"strings"
	"time"
)

type Register struct {
	DialTimeout time.Duration
	EtcdAddr    []string

	leaseID     clientv3.LeaseID
	keepAliveCh <-chan *clientv3.LeaseKeepAliveResponse
	closeCh     chan struct{}

	srvInfo *Server
	srvTTL  time.Duration
	client  *clientv3.Client
}

// NewRegister 新建 Register 对象
func NewRegister(etcdAddr []string) *Register {
	return &Register{
		EtcdAddr:    etcdAddr,
		DialTimeout: 3 * time.Second,
	}
}

// Start 开始注册服务
func (r *Register) Start(srvInfo *Server, ttl time.Duration) error {
	var err error

	if strings.Split(srvInfo.Addr, ":")[0] == "" {
		return errors.New("invalid ip addr")
	}

	if r.client, err = clientv3.New(clientv3.Config{
		Endpoints:   r.EtcdAddr,
		DialTimeout: r.DialTimeout,
	}); err != nil {
		return err
	}

	r.srvInfo = srvInfo
	r.srvTTL = ttl

	if err = r.register(); err != nil {
		return err
	}

	r.closeCh = make(chan struct{})
	go r.keepAlive()

	return nil

}

func (r *Register) register() error {
	ctx, cancel := context.WithTimeout(context.Background(), r.DialTimeout)
	defer cancel()

	leaseResp, err := r.client.Grant(ctx, int64(r.srvTTL/time.Second))
	if err != nil {
		return err
	}

	r.leaseID = leaseResp.ID

	if r.keepAliveCh, err = r.client.KeepAlive(context.Background(), r.leaseID); err != nil {
		return err
	}

	data, err := json.Marshal(r.srvInfo)
	if err != nil {
		return err
	}

	_, err = r.client.Put(context.Background(), BuildRegisterPath(r.srvInfo), string(data), clientv3.WithLease(r.leaseID))

	return err
}

// keepAlive 心跳保活
func (r *Register) keepAlive() {
	ticker := time.NewTicker(r.srvTTL)

	for {
		select {
		case <-r.closeCh:
			if err := r.unregister(); err != nil {
				logx.LogrusObj.Error("[DISCOVERY] unregister failed, error: ", err)
			}
			if _, err := r.client.Revoke(context.Background(), r.leaseID); err != nil {
				logx.LogrusObj.Error("[DISCOVERY] revoke failed, error: ", err)
			}
			logx.LogrusObj.Infoln("[DISCOVERY] unregister rpc server success")
			return
		case res := <-r.keepAliveCh:
			if res == nil {
				if err := r.register(); err != nil {
					logx.LogrusObj.Error("[DISCOVERY] register failed, error: ", err)
				}
			}
		case <-ticker.C:
			if r.keepAliveCh == nil {
				if err := r.register(); err != nil {
					logx.LogrusObj.Error("[DISCOVERY] register failed, error: ", err)
				}
			}
		}
	}
}

// unregister 注销服务
func (r *Register) unregister() error {
	_, err := r.client.Delete(context.Background(), BuildRegisterPath(r.srvInfo))
	return err
}

func (r *Register) Stop() {
	r.closeCh <- struct{}{}
}

// RpcEtcdRegister 在 etcd 中注册 rpc 服务，并保持心跳
func RpcEtcdRegister(topic string) *Register {
	register := NewRegister(viper.GetStringSlice("etcd.addr"))

	if err := register.Start(&Server{
		Name:   topic,
		Addr:   viper.GetString("server.ip") + ":" + viper.GetString("server.port"),
		Weight: viper.GetInt("server.weight"),
	}, 3*time.Second); err != nil {
		logx.LogrusObj.Fatalln(errors.Wrap(err, "[DISCOVERY]"))
	}

	return register
}
